package search.eightpuzzle.astar;

import search.eightpuzzle.model.PuzzleState;

import java.util.*;
import java.util.function.ToDoubleBiFunction;

public class AStar {

    /**
     * This is a mapping from a state (its unique hashcode) to the lowest cost we
     * found for this state. So it is the CLOSED list, but also stores the costs
     * for nodes on OPEN.
     */
    private final Map<Integer, Double> minCosts;
    private PuzzleState goal;
    private Heuristic heuristic;
    private Queue<WeightedState> open;
    private PuzzleState start;
    private int stepsCounter;

    public AStar(final PuzzleState start, final PuzzleState goal, final Heuristic heuristic) {
        assert start != null;
        assert goal != null;
        assert heuristic != null;

        this.start = start;
        this.goal = goal;
        this.heuristic = heuristic;
        this.open = new PriorityQueue<>();
        this.minCosts = new HashMap<>();
    }

    protected ToDoubleBiFunction<Double, PuzzleState> getEstimatedCostFunction() {
        return (currentCost, puzzleState) -> currentCost +  heuristic.estimate(puzzleState, goal);
    }

    protected Heuristic getHeuristic() {
        return heuristic;
    }

    protected PuzzleState getGoal() {
        return goal;
    }

    public PuzzleState search() {
        PuzzleState state = null;
        open.add(new WeightedState(0, start, getEstimatedCostFunction()));
        minCosts.put(start.hashCode(), 0.);
        stepsCounter = 0;
        while (state == null) {
            ++stepsCounter; // Number of expanded nodes
            try {
                state = this.advance(); // Expand next on open
            } catch (NoSuchElementException e) {
                // No node left on open
                return null;
            }
        }
        return state;
    }

    private PuzzleState advance() {
        final WeightedState n = this.nextStateFromOpen();
        if (n == null) {
            throw new NoSuchElementException(); // OPEN emtpy: FAIL
        }
//    System.out.println("Checking " + n.state + " with cost " + n.currentCost + " and total estimate " + n.estimatedCost);
        if (n.state.equals(this.goal)) {
            return n.state; // found goal
        }

        for (PuzzleState successor : n.state) { // Expand n
            double newCost = n.currentCost + 1; // We have always a cost of 1
            final Double oldCost = this.minCosts.get(successor.hashCode());
            if (oldCost == null || newCost < oldCost) {
                this.open.add(new WeightedState(newCost, successor, getEstimatedCostFunction()));
                this.minCosts.put(successor.hashCode(), newCost);
                // Note that I do not delete the old state from OPEN but then ignore
                // it later on. (in nextStateFromOpen()). Costs more memory, but safes
                // time.
            }
        }
        return null;
    }

    private WeightedState nextStateFromOpen() {
        WeightedState n = this.open.poll();
        if (n == null) {
            return null;
        }  // Open is empty
        while (n.currentCost > this.minCosts.get(n.state.hashCode())) {
            // A smaller path had been found and used...
            // Note: in A* this would have been removed from OPEN when the smaller
            // path was found. In my implementation it is not deleted, but ignored
            // here.
            n = this.open.poll();
            if (n == null) {
                return null;
            }
        }
        return n;
    }

    public int getStepsCounter() {
        return stepsCounter;
    }

    /**
     * A state including currentCost (g) and estimatedCost (f). The cost is
     * estimated in the constructor. The estimatedCost is used for sorting the
     * states in OPEN.
     */
    public static class WeightedState implements Comparable<WeightedState> {

        private final PuzzleState state;
        private double currentCost;
        private double estimatedCost;

        public WeightedState(double currentCost, PuzzleState state, ToDoubleBiFunction<Double, PuzzleState> estimatedCost) {
            assert state != null;
            this.currentCost = currentCost;
            this.state = state;
            this.estimatedCost = estimatedCost.applyAsDouble(currentCost, state);
        }

        @Override
        public int compareTo(WeightedState o) {
            // Compares this object with the specified object for order.
            // Returns a negative integer, zero, or a positive integer as this object
            // is less than, equal to, or greater than the specified object.
            final double weightDifference = this.estimatedCost - o.estimatedCost;
            return (int) weightDifference;
        }

    }

}
